package com.wangzhen.abstractroutingdatasource.config;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.util.concurrent.ThreadLocalRandom;

/**
 * 重写 determineCurrentLookupKey
 */
public class MyRoutingDataSource extends AbstractRoutingDataSource {

    @Autowired
    private DataSourceProperties dataSourceProperties;

    private final Logger logger = LoggerFactory.getLogger(MyRoutingDataSource.class);

    @Override
    protected Object determineCurrentLookupKey() {
        final String dataSourceType = DataSourceContextHolder.getDataSourceType();

        final int masterCount = dataSourceProperties.getMasters().size();
        final int slaveCount = dataSourceProperties.getSlaves().size();

        int index;
        if ( DataSourceContextHolder.WRITE.equalsIgnoreCase(dataSourceType)) {
            //使用随机数决定使用哪个写库
            index = ThreadLocalRandom.current().nextInt(masterCount) % masterCount;
            logger.info("使用了写库 {}", index);
        } else {
            //使用随机数决定使用哪个读库
            index = ThreadLocalRandom.current().nextInt(slaveCount) % slaveCount;
            logger.info("使用了读库 {}", index);
        }
        return dataSourceType + index;
    }
}

